{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This file serves the purpose of transforming raw interaction data into the data forms required to execute IDRec baselines.\n",
    "\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "SEQ_LEN = 10\n",
    "file_l = ['10wu_ks_pairs.tsv']\n",
    "data_l = ['ks']\n",
    "for idx in range(len(file_l)):\n",
    "    dat_seq = pd.read_csv(file_l[idx], sep='\\t',header=None)\n",
    "    dat_arr = np.array(dat_seq)\n",
    "    inter = []\n",
    "    for seq in dat_arr:\n",
    "        uid = seq[0]\n",
    "        iseq = seq[1].split()\n",
    "        for i, item in enumerate(iseq):\n",
    "            inter.append([item, uid, i])\n",
    "\n",
    "    inter_df = np.array(inter)\n",
    "    dat = pd.DataFrame(inter_df)\n",
    "    dat.columns = ['item_id', 'user_id', 'timestamp']\n",
    "    dat['timestamp'] = dat['timestamp'].astype(int)\n",
    "    dat.sort_values(by='timestamp', inplace=True, ascending=True)\n",
    "    user_list = dat['user_id'].values\n",
    "    item_list = dat['item_id'].values\n",
    "\n",
    "    index = {}\n",
    "    for i, key in enumerate(user_list):\n",
    "        if key not in index:\n",
    "            index[key] = [i]\n",
    "        else:\n",
    "            index[key].append(i)\n",
    "\n",
    "            indices = []\n",
    "        \n",
    "    for index in index.values():\n",
    "        indices.extend(list(index)[-(SEQ_LEN+3):])\n",
    "\n",
    "    final_dat = dict()\n",
    "    for k in dat:\n",
    "        final_dat[k] = dat[k].values[indices]\n",
    "    \n",
    "    final_dat = pd.DataFrame(final_dat)\n",
    "    print(final_dat.head(3))\n",
    "    print(final_dat['user_id'].nunique(),final_dat['item_id'].nunique(),final_dat.shape[0] )\n",
    "    os.makedirs(f'./{data_l[idx]}/', exist_ok=True)\n",
    "    final_dat.to_csv(f'./{data_l[idx]}/{data_l[idx]}.inter', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The following part generates the popularity count file (i.e. the pop.npy file needed in baseline code) of the dataset\n",
    "\n",
    "\n",
    "SEQ_LEN = 10\n",
    "class Data:\n",
    "    def __init__(self, df):\n",
    "        self.inter_feat = df\n",
    "        self._data_processing()\n",
    "\n",
    "\n",
    "    def _data_processing(self):\n",
    "        \n",
    "        self.id2token = {}\n",
    "        self.token2id = {}\n",
    "        remap_list = ['user_id', 'item_id']\n",
    "        for feature in remap_list:\n",
    "            feats = self.inter_feat[feature]\n",
    "            new_ids_list, mp = pd.factorize(feats)\n",
    "            mp = np.array(['[PAD]'] + list(mp))\n",
    "            token_id = {t: i for i, t in enumerate(mp)}\n",
    "            self.id2token[feature] = mp\n",
    "            self.token2id[feature] = token_id\n",
    "            self.inter_feat[feature] = new_ids_list+1\n",
    "        \n",
    "        self.user_num = len(self.id2token['user_id'])\n",
    "        self.item_num = len(self.id2token['item_id'])\n",
    "        self.inter_num = len(self.inter_feat)\n",
    "        self.uid_field = 'user_id'\n",
    "        self.iid_field = 'item_id'\n",
    "        self.user_seq = None\n",
    "        self.train_feat = None\n",
    "        self.feat_name_list = ['inter_feat']    \n",
    "\n",
    "   \n",
    "    def build(self):\n",
    "        \n",
    "        self.sort(by='timestamp')\n",
    "        user_list = self.inter_feat['user_id'].values\n",
    "        item_list = self.inter_feat['item_id'].values\n",
    "        grouped_index = self._grouped_index(user_list)\n",
    "        \n",
    "        user_seq = {}\n",
    "        for uid, index in grouped_index.items():\n",
    "            user_seq[uid] = item_list[index]\n",
    "    \n",
    "        self.user_seq = user_seq\n",
    "        train_feat = dict()\n",
    "        test_feat = dict()\n",
    "        valid_feat = dict()\n",
    "        indices = []\n",
    "       \n",
    "        for index in grouped_index.values():\n",
    "            indices.extend(list(index)[:-2])\n",
    "        for k in self.inter_feat:\n",
    "            train_feat[k] = self.inter_feat[k].values[indices]\n",
    "        \n",
    "        indices = []\n",
    "        for index in grouped_index.values():\n",
    "            indices.extend([index[-2]])\n",
    "        for k in self.inter_feat:\n",
    "            valid_feat[k] = self.inter_feat[k].values[indices]\n",
    "\n",
    "        indices = []\n",
    "        for index in grouped_index.values():\n",
    "            indices.extend([index[-1]])\n",
    "        for k in self.inter_feat:\n",
    "            test_feat[k] = self.inter_feat[k].values[indices]          \n",
    "                \n",
    "        self.train_feat = train_feat\n",
    "        return train_feat, valid_feat, test_feat\n",
    "    \n",
    "\n",
    "    def _grouped_index(self, group_by_list):\n",
    "        index = {}\n",
    "        for i, key in enumerate(group_by_list):\n",
    "            if key not in index:\n",
    "                index[key] = [i]\n",
    "            else:\n",
    "                index[key].append(i)\n",
    "        return index\n",
    "    \n",
    "    def _build_seq(self, train_feat):\n",
    "        max_item_list_len = SEQ_LEN+1               \n",
    "        uid_list, item_list_index= [], []\n",
    "        seq_start = 0\n",
    "        save = False\n",
    "        user_list = train_feat['user_id']\n",
    "        user_list = np.append(user_list, -1)\n",
    "        last_uid = user_list[0]\n",
    "        for i, uid in enumerate(user_list):\n",
    "            if last_uid != uid :\n",
    "                save = True \n",
    "            if save:\n",
    "                if i - seq_start > max_item_list_len:\n",
    "                    offset = (i - seq_start) % max_item_list_len\n",
    "                    seq_start += offset\n",
    "                    x = torch.arange(seq_start, i)\n",
    "                    sx = torch.split(x, max_item_list_len)\n",
    "                    for sub in sx:\n",
    "                        uid_list.append(last_uid)\n",
    "                        item_list_index.append(slice(sub[0],sub[-1]+1)) \n",
    "                                             \n",
    "                        \n",
    "                else:\n",
    "                    uid_list.append(last_uid)\n",
    "                    item_list_index.append(slice(seq_start,i))\n",
    "            \n",
    "                    \n",
    "                save = False\n",
    "                last_uid = uid\n",
    "                seq_start = i\n",
    "        \n",
    "        seq_train_feat = {}\n",
    "        seq_train_feat['user_id'] = np.array(uid_list)\n",
    "        seq_train_feat['item_seq'] = []\n",
    "        seq_train_item = []\n",
    "        for index in item_list_index:\n",
    "            seq_train_feat['item_seq'].append(train_feat['item_id'][index])\n",
    "            seq_train_item+=list(train_feat['item_id'][index])\n",
    "        \n",
    "        self.seq_train_item = seq_train_item\n",
    "        return seq_train_feat\n",
    "\n",
    "\n",
    "    def sort(self, by, ascending=True):\n",
    "        self.inter_feat.sort_values(by=by, ascending=ascending, inplace=True)\n",
    "        \n",
    "    \n",
    "\n",
    "data_list = ['ks', ]\n",
    "\n",
    "for idx in range(len(data_list)):\n",
    "    inter = pd.read_csv(f'./{data_list[idx]}/{data_list[idx]}.inter', delimiter=',', dtype={'item_id':str, 'user_id':str, 'timestamp':int}, header=0, names=['item_id', 'user_id', 'timestamp'] \n",
    "            )\n",
    "\n",
    "    item_num = inter['item_id'].nunique()\n",
    "    D  = Data(inter)\n",
    "    train, valid, test = D.build()\n",
    "    D._build_seq(train)\n",
    "    train_items = D.seq_train_item\n",
    "    train_item_counts = [0] * (item_num + 1)\n",
    "    for i in train_items:\n",
    "        train_item_counts[i] += 1\n",
    "    item_counts_powered = np.power(train_item_counts, 1.0)\n",
    "    pop_prob_list = []\n",
    "\n",
    "    for i in range(1, item_num + 1):\n",
    "        pop_prob_list.append(item_counts_powered[i])\n",
    "    pop_prob_list = pop_prob_list / sum(np.array(pop_prob_list))\n",
    "    pop_prob_list = np.append([1], pop_prob_list)\n",
    "    print(('prob max: {}, prob min: {}, prob mean: {}'.\\\n",
    "            format(max(pop_prob_list), min(pop_prob_list), np.mean(pop_prob_list))))\n",
    "\n",
    "    np.save(f'./{data_list[idx]}/pop',pop_prob_list)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
